package com.sdy.common.utils;

import java.util.ArrayList;
import java.util.Collection;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.function.Function;

public class MapUtil {
    public static <PK, T> Map<PK, T> collectionToMap(Collection<T> collection, PkAccess<PK, T> pkAccess) {
        Map<PK, T> map = new HashMap<>(32);
        if (collection != null) {
            collection.forEach(item -> map.put(pkAccess.fetch(item), item));
        }
        return map;
    }
    
    public static <PK, T> Map<PK, List<T>> collectionToMapList(Collection<T> collection, PkAccess<PK, T> pkAccess) {
        Map<PK, List<T>> map = new HashMap<>(32);
        if (collection != null) {
            collection.forEach(item -> {
                PK key = pkAccess.fetch(item);
                map.computeIfAbsent(key, k -> new ArrayList<>());
                map.get(key).add(item);
            });
        }
        return map;
    }

    /**
     * 遍历树
     */
    public static <T, R> void traverseTree(T tree, Function<T, Collection<T>> childrenFunction, Function<T, R> doFunction) {
        if (tree == null) {
            return;
        }
        doFunction.apply(tree);
        Collection<T> children = childrenFunction.apply(tree);
        if (children != null) {
            children.forEach(child -> traverseTree(child, childrenFunction, doFunction));
        }
    }

    /**
     * 遍历树
     */
    public static <T, R> void traverseTree(Collection<T> tree, Function<T, Collection<T>> childrenFunction, Function<T, R> doFunction) {
        if (tree == null || tree.isEmpty()) {
            return;
        }
        tree.forEach(child -> traverseTree(child, childrenFunction, doFunction));
    }
    
    @FunctionalInterface
    public interface PkAccess<PK, T> {
        PK fetch(T t);
    }
}
